# -*- coding: utf-8 -*
from flyai.framework import FlyAI
from net import get_model
from torchvision import transforms
import torch
import os
from path import MODEL_PATH
from PIL import Image
from efficientdet.dataset import Resizer, Normalizer, Augmenter, collater
from utils.utils import get_last_weights
from backbone import EfficientDetBackbone
import cv2
import numpy as np
from efficientdet.utils import BBoxTransform, ClipBoxes
from utils.utils import preprocess, invert_affine, postprocess, STANDARD_COLORS, standard_to_bgr, get_index_label, plot_one_box
from torch.backends import cudnn
from net import get_model_efficientdet


# 判断gpu是否可用
if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

compound_coef = 3
force_input_size = None  # set None to use default size
use_cuda =True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
threshold = 0.2
iou_threshold = 0.2

color_list = standard_to_bgr(STANDARD_COLORS)
obj_list = ['Tuberculosis']

input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size

#transform = transforms.Compose([
#    Normalizer(mean=[0,0,0],std=[0.7,0.7,0.7]),
#    Augmenter(),
#    Resizer(input_sizes[compound_coef])
#    ])


def get_return(boxes, labels, scores, image_name, label_id_name):
    ''' 输入示例：
    :param boxes: [[735.097,923.59283,770.1911,998.49335], [525.39496,535.89667,578.4822,589.8431 ]]]   box格式[Xmin, Ymin, Xmax, Ymax]
    :param labels: [1,1]
    :param scores: [0.2， 0.3]
    :param image_name: 0.jpg
    :param label_id_name: {1: 'TBbacillus'}  label id 到 label name的映射

    :return:  [{"image_name": '0.jpg', "label_name": 'TBbacillus', "bbox": [735, 923, 770-735, 998-923], "confidence": 0.2},
                {"image_name": '0.jpg', "label_name": 'TBbacillus', "bbox": [525, 535, 578-525, 589-535], "confidence": 0.3}]
                返回 box的格式为[xmin, ymin, width, height]
    最终评估方式采用coco数据集的map，具体可参考 https://cocodataset.org/
    '''
    result = []
    if len(boxes) == len(labels) == len(scores):
        for i in range(len(boxes)):
            box = boxes[i] # [Xmin, Ymin, Xmax, Ymax]
            bbox = [int(box[0]), int(box[1]), int(box[2]-box[0]), int(box[3]-box[1])] # [xmin, ymin, width, height]
            label_name = label_id_name[labels[i]]
            confidence = scores[i]
            ann = {"image_name": image_name, "label_name": label_name, "bbox": bbox, "confidence": confidence}
            result.append(ann)
    return result



class Prediction(FlyAI):
    def load_model(self):
        '''
        模型初始化，必须在此方法中加载模型
        '''
        anchor_ratios = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)]
        anchor_scales = [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]
        self.model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list),
                             ratios=anchor_ratios, scales=anchor_scales)
        print('load from best.pth...')
        model_file = get_last_weights(MODEL_PATH +'/Tuberculosis')
        print(model_file)
        self.model.load_state_dict(torch.load(model_file))
        self.model.requires_grad_(False)
        self.model.eval()
        if use_cuda:
            self.model = self.model.cuda()
        if use_float16:
            self.model = self.model.half()
        print('load model done...')
    
    def display(self,preds, imgs, imshow=True, imwrite=False):
        for i in range(len(imgs)):
            if len(preds[i]['rois']) == 0:
                continue

            for j in range(len(preds[i]['rois'])):
                x1, y1, x2, y2 = preds[i]['rois'][j].astype(np.int)
                #obj = obj_list[preds[i]['class_ids'][j]]
                score = float(preds[i]['scores'][j])
                #plot_one_box(imgs[i], [x1, y1, x2, y2], label=obj,score=score,color=color_list[get_index_label(obj, obj_list)])
                plot_one_box(imgs[i], [x1, y1, x2, y2], label="Tuberculosis",score=score,color='r')


            if imshow:
                cv2.imshow('img', imgs[i])
                cv2.waitKey(0)

            if imwrite:
                cv2.imwrite(f'test/img_inferred_d{compound_coef}_this_repo_{i}.jpg', imgs[i])

    def predict(self, image_path):
        '''
        模型预测返回结果
        :param input: 评估传入样例 {"image_path": "./data/input/image/0.jpg"}
        :return: 具体见 get_return 方法
        '''
        image_name = image_path.split('/')[-1]
        #img = Image.open(image_path)
        #img = cv2.imread(image_path)
        #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        #img.astype(np.float32) / 255.
        #img = transform(img)
        #img = img.unsqueeze(0)
        #img = img.to(device)

        self.model.eval()
        ori_imgs, framed_imgs, framed_metas = preprocess(image_path, max_size=input_size)

        if use_cuda:
            x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
        else:
            x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)

        x = x.to(torch.float32 if not use_float16 else torch.float16).permute(0, 3, 1, 2)

        print('*'*100)
        
        with torch.no_grad():
            features, regression, classification, anchors = self.model(x)

            regressBoxes = BBoxTransform()
            clipBoxes = ClipBoxes()

            out = postprocess(x,
                      anchors, regression, classification,
                      regressBoxes, clipBoxes,
                      threshold, iou_threshold)
        
        preds = invert_affine(framed_metas, out)
        preds = preds[0]
        #self.display(out, ori_imgs, imshow=True, imwrite=True)
        #preds = self.model(img)
        #preds = [{k: v.to(torch.device('cpu')) for k, v in t.items()} for t in preds]
        #preds = [{k: v.data.cpu().numpy() for k, v in t.items()} for t in preds][0]
        #print(preds)
        result = get_return(preds['rois'], preds['class_ids'], preds['scores'], image_name, label_id_name={0: 'TBbacillus'})
        print(result)
        return result

if __name__ == "__main__":
    pre = Prediction()
    pre.load_model()
    #result = pre.predict(image_path="./data/input/TBDetection/image/2702.jpg")
    pre.predict(image_path="./data/input/TBDetection/image/836.jpg")
    #print(len(result))